#ifndef THREAD_POOL_H
#define THREAD_POOL_H

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <pthread.h>

#define LIST_INSERT(item, list)  \
	do                           \
	{                            \
		item->prev = NULL;       \
		item->next = list;       \
		if ((list) != NULL)      \
			(list)->prev = item; \
		(list) = item;           \
	} while (0)

#define LIST_REMOVE(item, list)            \
	do                                     \
	{                                      \
		if (item->prev != NULL)            \
			item->prev->next = item->next; \
		if (item->next != NULL)            \
			item->next->prev = item->prev; \
		if (list == item)                  \
			list = item->next;             \
		item->prev = item->next = NULL;    \
	} while (0)

struct nTask
{
	void (*task_func)(void *task);
	void *user_data;
	void (*cleanup_func)(void *task); // 任务清理函数
	struct nTask *prev;
	struct nTask *next;
};

struct nWorker
{
	pthread_t threadid;
	int terminate;
	struct nManager *manager;

	struct nWorker *prev;
	struct nWorker *next;
};

typedef struct nManager
{
	struct nTask *tasks;
	struct nWorker *workers;
	int worker_count;
	int is_shutdown;
	pthread_mutex_t mutex;
	pthread_cond_t cond;
} ThreadPool;

// callback != task
static void *nThreadPoolCallback(void *arg)
{

	struct nWorker *worker = (struct nWorker *)arg;

	while (1)
	{

		pthread_mutex_lock(&worker->manager->mutex);
		while (worker->manager->tasks == NULL)
		{
			if (worker->terminate)
				break;
			pthread_cond_wait(&worker->manager->cond, &worker->manager->mutex);
		}
		if (worker->terminate)
		{
			pthread_mutex_unlock(&worker->manager->mutex);
			break;
		}

		struct nTask *task = worker->manager->tasks;
		LIST_REMOVE(task, worker->manager->tasks);

		pthread_mutex_unlock(&worker->manager->mutex);

		task->task_func(task);

		if (task == NULL)
			continue;

		if (task->cleanup_func)
			task->cleanup_func(task); // 自定义清理函数

		else
		{
			if (task->user_data)
			{
				free(task->user_data);
			}

			free(task);
		}
	}
	pthread_mutex_lock(&worker->manager->mutex);
	worker->manager->worker_count--;
	pthread_cond_signal(&worker->manager->cond);
	pthread_mutex_unlock(&worker->manager->mutex);

	free(worker);
}

// API
int nThreadPoolCreate(ThreadPool *pool, int numWorkers)
{

	if (pool == NULL)
		return -1;
	if (numWorkers < 1)
		numWorkers = 1;
	memset(pool, 0, sizeof(ThreadPool));

	pthread_cond_t blank_cond = PTHREAD_COND_INITIALIZER;
	memcpy(&pool->cond, &blank_cond, sizeof(pthread_cond_t));

	// pthread_mutex_init(&pool->mutex, NULL);
	pthread_mutex_t blank_mutex = PTHREAD_MUTEX_INITIALIZER;
	memcpy(&pool->mutex, &blank_mutex, sizeof(pthread_mutex_t));

	int i = 0;
	for (i = 0; i < numWorkers; i++)
	{
		struct nWorker *worker = (struct nWorker *)malloc(sizeof(struct nWorker));
		if (worker == NULL)
		{
			perror("malloc");
			return -2;
		}
		memset(worker, 0, sizeof(struct nWorker));
		worker->manager = pool; //

		int ret = pthread_create(&worker->threadid, NULL, nThreadPoolCallback, worker);
		if (ret)
		{
			perror("pthread_create");
			free(worker);
			return -3;
		}

		LIST_INSERT(worker, pool->workers);
		pool->worker_count++;
	}

	// success
	return 0;
}

// API
int nThreadPoolDestory(ThreadPool *pool)
{
	if (pool == NULL)
	{
		return -1;
	}

	// 如果已经关闭，直接返回
	if (pool->is_shutdown)
	{

		return 0;
	}

	pthread_mutex_lock(&pool->mutex);
	// 标记线程池为已关闭
	pool->is_shutdown = 1;

	// 标记所有工作线程为终止状态
	struct nWorker *worker = NULL;
	for (worker = pool->workers; worker != NULL; worker = worker->next)
	{
		worker->terminate = 1;
	}

	// 唤醒所有等待的工作线程
	pthread_cond_broadcast(&pool->cond);

	// 等待所有工作线程结束
	while (pool->worker_count > 0)
	{
		pthread_cond_wait(&pool->cond, &pool->mutex);
	}

	// 清理剩余任务
	struct nTask *task = NULL, *next_task = NULL;
	for (task = pool->tasks; task != NULL; task = next_task)
	{
		next_task = task->next;
		if (task->cleanup_func)
			task->cleanup_func(task); // 自定义清理函数
		else
		{
			if (task->user_data)
			{
				free(task->user_data);
			}

			free(task);
		}
	}

	// 清理互斥锁和条件变量
	pthread_mutex_unlock(&pool->mutex);
	pthread_mutex_destroy(&pool->mutex);
	pthread_cond_destroy(&pool->cond);

	// 清空指针
	pool->workers = NULL;
	pool->tasks = NULL;

	return 0;
}

// API
int nThreadPoolPushTask(ThreadPool *pool, struct nTask *task)
{

	pthread_mutex_lock(&pool->mutex);

	LIST_INSERT(task, pool->tasks);

	pthread_cond_signal(&pool->cond);

	pthread_mutex_unlock(&pool->mutex);
}

#endif // THREAD_POOL_H